import numpy as np
import cv2
import math

def trans_n8(mask, px, py):
    M = np.array([[1,0,py],[0,1,px]]).astype(np.float32)
    trans = cv2.warpAffine(mask, M, mask.shape[::-1])
    return trans

def flux_to_skl(flux, threshold, dks, eks):
    magnitude, angle = cv2.cartToPolar(flux[0][0], flux[0][1])
    vmax = float(magnitude.max())
    mask = (magnitude > threshold).astype(np.uint8)
    mask_rev = (mask == 0).astype(np.uint8)
    ending = trans_n8(mask_rev, -1, -1)*mask*np.logical_and(angle >= math.pi/8, angle < 3*math.pi/8) \
           + trans_n8(mask_rev, 0, -1)*mask*np.logical_and(angle >= 3*math.pi/8, angle < 5*math.pi/8) \
           + trans_n8(mask_rev, 1, -1)*mask*np.logical_and(angle >= 5*math.pi/8, angle < 7*math.pi/8) \
           + trans_n8(mask_rev, 1, 0)*mask*np.logical_and(angle >= 7*math.pi/8, angle < 9*math.pi/8) \
           + trans_n8(mask_rev, 1, 1)*mask*np.logical_and(angle >= 9*math.pi/8, angle < 11*math.pi/8) \
           + trans_n8(mask_rev, 0, 1)*mask*np.logical_and(angle >= 11*math.pi/8, angle < 13*math.pi/8) \
           + trans_n8(mask_rev, -1, 1)*mask*np.logical_and(angle >= 13*math.pi/8, angle < 15*math.pi/8) \
           + trans_n8(mask_rev, -1, 0)*mask*np.logical_or(angle < math.pi/8, angle >= 15*math.pi/8)
    ending = (ending > 0).astype(np.uint8)
    element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dks, dks))
    ending = cv2.dilate(ending, element)
    element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (eks, eks))
    ending = cv2.erode(ending, element)
    skl = ending * (vmax - magnitude) / vmax
    return skl